import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import butter, filtfilt, find_peaks, cheby2
from scipy.fft import fft, fftfreq
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

# Folder containing all subject files
folder_path = 'BIDMC_dataset/data/'  # Replace with the folder path containing your files

# Function to design a Butterworth low-pass filter
def butter_lowpass(cutoff, fs, order=4):
    nyquist = 0.5 * fs
    normal_cutoff = cutoff / nyquist
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    return b, a

# Function to apply the low-pass filter
def apply_filter(signal, cutoff, fs, order=4):
    b, a = butter_lowpass(cutoff, fs, order=order)
    return filtfilt(b, a, signal)

# Parameters
sampling_rate = 125  # Hz
cutoff_freq = 0.5  # Hz
window_durations = [10, 20, 30, 32, 64] + list(range(40, 130, 10))  # Added 32s and 64s

# Define output folder for saving CSV files
output_folder = 'Window_Size_Results'
os.makedirs(output_folder, exist_ok=True)

# Initialize dictionaries to store per-file RMSE, MAE, and R2
rmse_files = {wd: [] for wd in window_durations}
mae_files = {wd: [] for wd in window_durations}
r2_files = {wd: [] for wd in window_durations}

# Function to save results to CSV
def save_results_to_csv(window_size, subject_results, metric_name):
    file_path = os.path.join(output_folder, f'Window_{window_size}s_{metric_name}.csv')
    df = pd.DataFrame(subject_results, columns=['FileName', metric_name])
    df.to_csv(file_path, index=False)
    print(f"Saved {metric_name} results for window size {window_size}s to {file_path}")

# Loop over each subject file
for file_name in os.listdir(folder_path):
    if file_name.endswith('_Signals.csv'):
        file_identifier = file_name.replace('_Signals.csv', '')

        signals_file_path = os.path.join(folder_path, file_name)
        numerics_file_path = signals_file_path.replace('_Signals.csv', '_Numerics.csv')

        signals_data = pd.read_csv(signals_file_path)
        numerics_data = pd.read_csv(numerics_file_path)

        ppg_signal = signals_data[' PLETH'].values
        time_signals = signals_data['Time [s]'].values

        # 1. Signal normalization
        ppg_signal = (ppg_signal - np.mean(ppg_signal)) / np.std(ppg_signal)

        # 2. Chebyshev Type II bandpass filter (0.1 - 0.5 Hz)
        low = 0.1 / (0.5 * sampling_rate)
        high = 0.5 / (0.5 * sampling_rate)
        b_cheby, a_cheby = cheby2(N=4, rs=20, Wn=[low, high], btype='bandpass')
        filtered_signal = filtfilt(b_cheby, a_cheby, ppg_signal)

        # 3. IQR-based outlier removal
        Q1, Q3 = np.percentile(filtered_signal, [25, 75])
        IQR = Q3 - Q1
        lower_bound = Q1 - 1.5 * IQR
        upper_bound = Q3 + 1.5 * IQR
        filtered_signal = np.clip(filtered_signal, lower_bound, upper_bound)

        # 4. KSQI calculation
        ksqi = np.mean(((filtered_signal - np.mean(filtered_signal)) / np.std(filtered_signal))**4)
        print(f"{file_identifier} - KSQI: {ksqi:.2f}")

        # 5. FFT Analysis (optional visualization)
        fft_vals = fft(filtered_signal)
        fft_freqs = fftfreq(len(filtered_signal), d=1/sampling_rate)

        # 6. Butterworth low-pass filter
        filtered_signal = apply_filter(filtered_signal, cutoff=cutoff_freq, fs=sampling_rate)

        # Extract the actual respiration rate
        resp_actual = numerics_data[' RESP']
        time_numerics = numerics_data['Time [s]']

        # 7 & 8: Peak detection and RR estimation
        for window_duration in window_durations:
            respiration_rates_calculated = []
            respiration_rates_actual = []

            for start_time in range(0, 121, window_duration):
                end_time = start_time + window_duration

                segment_mask_signals = (time_signals >= start_time) & (time_signals < end_time)
                segment_signal = filtered_signal[segment_mask_signals]

                peaks, _ = find_peaks(segment_signal, distance=sampling_rate)

                if len(peaks) > 1:
                    resp_intervals = np.diff(peaks) / sampling_rate
                    respiration_rate_bpm = 60 / np.mean(resp_intervals)
                else:
                    respiration_rate_bpm = np.nan
                respiration_rates_calculated.append(respiration_rate_bpm)

                segment_mask_numerics = (time_numerics >= start_time) & (time_numerics < end_time)
                segment_resp_actual = resp_actual[segment_mask_numerics]

                if len(segment_resp_actual) > 0:
                    respiration_rate_actual = segment_resp_actual.mean()
                else:
                    respiration_rate_actual = np.nan
                respiration_rates_actual.append(respiration_rate_actual)

            valid_indices = np.isfinite(respiration_rates_calculated) & np.isfinite(respiration_rates_actual)
            respiration_rates_calculated = np.array(respiration_rates_calculated)[valid_indices]
            respiration_rates_actual = np.array(respiration_rates_actual)[valid_indices]

            if len(respiration_rates_actual) > 1 and len(respiration_rates_calculated) > 1:
                avg_rmse = np.sqrt(mean_squared_error(respiration_rates_actual, respiration_rates_calculated))
                avg_mae = mean_absolute_error(respiration_rates_actual, respiration_rates_calculated)
                avg_r2 = r2_score(respiration_rates_actual, respiration_rates_calculated) if np.var(respiration_rates_actual) > 0 else np.nan
            else:
                avg_rmse = avg_mae = avg_r2 = np.nan

            rmse_files[window_duration].append((file_identifier, avg_rmse))
            mae_files[window_duration].append((file_identifier, avg_mae))
            r2_files[window_duration].append((file_identifier, avg_r2))

# Save results
for window_duration in window_durations:
    save_results_to_csv(window_duration, rmse_files[window_duration], 'RMSE')
    save_results_to_csv(window_duration, mae_files[window_duration], 'MAE')
    save_results_to_csv(window_duration, r2_files[window_duration], 'R2')

print("\nResults saved for all window sizes.")